package grpc

import (
	"context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/connectivity"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/keepalive"
	"net"
	"sync"
	"time"
)

var (
	ConnPool = make(map[string]*grpc.ClientConn) // 连接池
	poolMux  = &sync.RWMutex{}                   // 连接池锁
)

// GetClientConn 获取agent连接（自动初始化/重建）
func GetClientConn(agentAddr string) (*grpc.ClientConn, error) {
	// 第一层检查：快速读取
	poolMux.RLock()
	conn, exists := ConnPool[agentAddr]
	poolMux.RUnlock()

	if exists {
		state := conn.GetState()
		if state == connectivity.Ready || state == connectivity.Idle {
			return conn, nil
		}
		// 仅在连接已关闭时才主动关闭（Shutdown是终态）
		if state == connectivity.Shutdown {
			_ = conn.Close()
		}
		// TransientFailure/Connecting 状态：不关闭，等待gRPC自动处理
	}

	// 第二层检查：加锁创建
	poolMux.Lock()
	defer poolMux.Unlock()

	// 再次检查是否已被其他协程创建
	if conn, exists := ConnPool[agentAddr]; exists {
		if state := conn.GetState(); state == connectivity.Ready || state == connectivity.Idle {
			return conn, nil
		}
		_ = conn.Close()
	}

	// 创建新连接
	kp := keepalive.ClientParameters{
		Time:                15 * time.Second, // 每15秒发送一次ping（建议10~30秒）
		Timeout:             5 * time.Second,  // ping超时时间（建议小于Time）
		PermitWithoutStream: true,             // 关键：无活跃流时也发送心跳
	}

	conn, err := grpc.NewClient(
		agentAddr,
		grpc.WithTransportCredentials(insecure.NewCredentials()),
		grpc.WithKeepaliveParams(kp),
		grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) {
			return net.DialTimeout("tcp", addr, 5*time.Second) // 同步拨号超时
		}),
	)
	if err != nil {
		return nil, err
	}
	ConnPool[agentAddr] = conn
	return conn, nil
}

// CloseAllConnections 关闭所有连接（优雅退出时调用）
func CloseAllConnections() {
	poolMux.Lock()
	defer poolMux.Unlock()

	for addr, conn := range ConnPool {
		_ = conn.Close()
		delete(ConnPool, addr)
	}
}
